import torch
from torch import nn

from util import try_gpu
from dataload import get_dataload

class Evaluator:
    def __init__(self, batch_size, is_train):
        self.dataloader = get_dataload("./BigData/dataset", my_batch_size=batch_size, is_train=is_train)
    
    def eval(self, model):
        model.eval()
        cnt = 0
        cnt_true = 0
        for input, label in self.dataloader:
            input = input.to(device=try_gpu())
            label = label.to(device=try_gpu())
            output = model(input)
            pred = torch.softmax(output, dim=1)
            pred = torch.argmax(output, dim=1)
            for i in range(len(label)):
                cnt += 1
                if pred[i] == label[i]:
                    cnt_true += 1
        accuracy = (float)(cnt_true / cnt)
        return accuracy